import requests, os, yaml
from lxml import etree

headers={ 'User-Agent':'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/72.0.3626.121 Safari/537.36', }

def _download_icon(icon_url, filepath):
    resp = requests.get(icon_url, headers=headers, timeout=3)
    # if (resp.status_code == 200 and len(resp.history)==0) or (len(resp.history)!=0 and resp.history[0].status_code==301):
    try: 
        resp.content.decode('utf-8')
        binary = False
    except: binary = True
    if resp.status_code == 200 and binary:
        with open(filepath, 'wb') as f: f.write(resp.content)
        print(f'download {icon_url} to {filepath}'); return True
    else: print(f'{filepath} icon doesnt exist, search icon_url'); return False

def _search_icon_url(url):
    resp = requests.get(url, headers=headers, timeout=5)
    html = etree.HTML(resp.text)
    ls = html.xpath('//@href')
    
    ls = list(filter(lambda x: x.find('favicon')>=0 and (x.endswith('png') or x.endswith('ico') or x.find('?')>=0), ls))

    if len(ls) > 0:
        icon_url = str(ls[0])
        if icon_url.startswith('//'): icon_url = 'https:' + icon_url
        elif icon_url.startswith('/'): icon_url = url + icon_url
        return icon_url
    else: return ""

def download_icon(url, filepath):
    get_base_url = lambda x: x.replace('//','||').split('/')[0].replace('||','//')
    icon_url = get_base_url(url) + '/favicon.ico'

    if os.path.exists(filepath):  return

    try:
        if not _download_icon(icon_url, filepath):
            icon_url = _search_icon_url(get_base_url(url))
            if not icon_url: print('{0} search failed'.format(filepath))
            else: _download_icon(icon_url, filepath)
    except: 
        print('download timeout')


with open('navi.yaml', 'r', encoding='utf-8') as f:
    obj = yaml.load(f, yaml.FullLoader)
    for k, v in obj.items():
        for item in v:
            download_icon(item['link'], './img/{0}.ico'.format(item['name']))
        
